package cn.edu.flink.scala.tutorial.topn

import cn.edu.flink.scala.tutorial.source.ClickSource
import org.apache.flink.api.common.functions.AggregateFunction
import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows
import org.apache.flink.streaming.api.windowing.time.Time
import org.apache.flink.streaming.api.windowing.windows.TimeWindow
import org.apache.flink.util.Collector

import scala.collection.convert.ImplicitConversions.`iterable AsScalaIterable`


object KeyedProcessFunctionTest {
  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()

    val urlStream = env.addSource(new ClickSource).assignAscendingTimestamps(_.timestamp).map(_.url)

    val urlCountStream = urlStream.keyBy(_.toString)
      .window(TumblingEventTimeWindows.of(Time.seconds(10L)))
      .aggregate(
        new AggregateFunction[String, Long, Long] {
          override def createAccumulator(): Long = 0L

          override def add(value: String, accumulator: Long): Long = accumulator + 1L

          override def getResult(accumulator: Long): Long = accumulator

          override def merge(a: Long, b: Long): Long = ???
        },
        new ProcessWindowFunction[Long, UrlViewCount, String, TimeWindow] {
          override def process(key: String, context: Context, elements: Iterable[Long], out: Collector[UrlViewCount]): Unit = {
            out.collect(UrlViewCount(key, elements.iterator.next(), context.window.getStart, context.window.getEnd))
          }
        })

    urlCountStream.keyBy(_.windowEnd)
      .process(new KeyedProcessFunction[Long, UrlViewCount, String]{
        var urlViewCountListState: ListState[UrlViewCount] = _


        override def open(parameters: Configuration): Unit = {
          urlViewCountListState = getRuntimeContext.getListState(new ListStateDescriptor[UrlViewCount]("list-state", classOf[UrlViewCount]))
        }

        override def processElement(value: UrlViewCount, ctx: KeyedProcessFunction[Long, UrlViewCount, String]#Context, out: Collector[String]): Unit = {
          urlViewCountListState.add(value)

          // 注册定时器
          ctx.timerService().registerEventTimeTimer(value.windowEnd + 1)
        }

        override def onTimer(timestamp: Long, ctx: KeyedProcessFunction[Long, UrlViewCount, String]#OnTimerContext, out: Collector[String]): Unit = {

          val urlViewCountList = urlViewCountListState.get().toList
          urlViewCountListState.clear()
          val topn = urlViewCountList.sortBy(-_.count).take(3)

          val builder = new StringBuilder()
          builder.append(s"==================窗口: ${topn.head.windowStart}~${topn.head.windowEnd}====================\n")
          for (i <- topn.indices){
            builder.append(s"Top${i+1} url: ${topn(i).url}, count: ${topn(i).count} \n")
          }

          builder.append("\n")
          out.collect(builder.toString())
        }
      })
      .print()

    env.execute("KeyedProcessFunctionTest")
  }

  case class UrlViewCount(url: String, count: Long, windowStart: Long, windowEnd: Long)
}
